#include "pch.h"
#include "DeviceResources.h"
#include "Camera.h"
#include "ParticleSystem.hpp"
#include "ForwardingVS.hpp"
#include "Generated/ParticleExpansionGS.hpp"
#include "Generated/ParticleGenerationGS.hpp"
#include "Generated/ParticleRenderingPS.hpp"
#include "Generated/ParticleRenderingVS.hpp"

using namespace DirectX;
using namespace DX;

struct ParticleGenGSCbPerFrame
{
    XMFLOAT3 initialWorldPos;
    float timeElapsed;
    float gameTime;
    float padding[3];
};

struct ParticleGenGSCbFixed
{
    float startLifetime;
    float padding[3];
};

struct alignas(16) ParticleExpansionGSCbView
{
    XMMATRIX viewProjMatrix;
    XMFLOAT3 eyePos;
};

struct ParticleRenderingVSCbFixed
{
    XMFLOAT3 accel;
    float startLifetime;
};

struct Particle
{
    DirectX::XMFLOAT4 initialPos;
    DirectX::XMFLOAT4 intialVel;
    DirectX::XMFLOAT2 size;
    //the last two should be packed 
    float currentAge;
    std::uint32_t type;
};


void ParticleSystem::Draw(DX::DeviceResources& deviceResources, ID3D11DeviceContext& context3D, const Camera& camera)
{
    deviceResources.PIXBeginEvent(L"Particle generation.");
    ID3D11Buffer* inputBuffer = nullptr;
    if (m_firstRun)
    {
        inputBuffer = m_initialParticle.Get();
    }
    else
    {
        inputBuffer = m_particleBuffer1.Get();
    }
    this->UpdateConstantBuffers(context3D, camera);
    context3D.PSSetShader(nullptr, nullptr, 0);
    context3D.VSSetShader(m_forwardingVS.Get(), nullptr, 0);
    const UINT particleStride = static_cast<UINT>(sizeof(Particle));
    const UINT particleOffset = 0;
    context3D.IASetVertexBuffers(0, 1, &inputBuffer, &particleStride,
                                 &particleOffset);
    context3D.IASetInputLayout(m_particleInputLayout.Get());
    context3D.IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_POINTLIST);
    SetupGenerationGS(context3D);
    ID3D11Buffer* gsOutputBuffers[] = { m_particleBuffer2.Get() };
    const UINT gsOutputBufferOffsets[] = {0};
    context3D.SOSetTargets(1, gsOutputBuffers, gsOutputBufferOffsets);
    ID3D11RenderTargetView* mainRT;
    ID3D11DepthStencilView* mainDsv;
    context3D.OMGetRenderTargets(1, &mainRT, &mainDsv);
    ID3D11RenderTargetView* emptyRT = nullptr;
    context3D.OMSetRenderTargets(1, &emptyRT, nullptr);
    context3D.OMSetDepthStencilState(m_disableDepthStencilState.Get(), 0);
    if (m_firstRun)
    {
        context3D.Draw(1, 0);
    }
    else
    {
        context3D.DrawAuto();
    }
    deviceResources.PIXEndEvent();

    deviceResources.PIXBeginEvent(L"Particle expandation.");

    //particles have been generated into m_particleBuffer2
    gsOutputBuffers[0] = nullptr;
    context3D.SOSetTargets(1, gsOutputBuffers, gsOutputBufferOffsets);

    ID3D11Buffer* bufferForRendering = m_particleBuffer2.Get();
    context3D.IASetVertexBuffers(0, 1, &bufferForRendering, &particleStride, &particleOffset);
    context3D.OMSetDepthStencilState(nullptr, 0);
    SetupExpansionGS(context3D);
    SetupRenderingPS(context3D);
    SetupRenderingVS(context3D);
    context3D.OMSetBlendState(m_alphaBlending.Get(), nullptr, 0xffffffff);
    context3D.OMSetRenderTargets(1, &mainRT, mainDsv);

    context3D.DrawAuto();
    using std::swap;
    swap(m_particleBuffer1, m_particleBuffer2);
    m_firstRun = false;

    deviceResources.PIXEndEvent();
}

void ParticleSystem::SetBumpInterval(float interval)
{
    m_bumpInterval = interval;
}

void ParticleSystem::CreateRandomTexture(ID3D11Device& device)
{
    const int kRandomNumCount = 1024;
    CD3D11_TEXTURE1D_DESC randomTextureDesc{DXGI_FORMAT_R32G32B32A32_FLOAT,
                                            kRandomNumCount, 1, 1};
    std::array<DirectX::XMFLOAT4, kRandomNumCount> cpuRandom;
    std::random_device rd;
    std::mt19937 g(rd());
    std::uniform_real_distribution<float> distribution{0.0f, 1.0f};
    auto genFloatRandom = [&]() { return distribution(g); };
    std::generate(cpuRandom.begin(), cpuRandom.end(), [&]() {
        return DirectX::XMFLOAT4{genFloatRandom(), genFloatRandom(),
                                 genFloatRandom(), genFloatRandom()};
    });
    D3D11_SUBRESOURCE_DATA randomTextureData{.pSysMem = cpuRandom.data()};
    DX::ThrowIfFailed(device.CreateTexture1D(&randomTextureDesc,
                                             &randomTextureData,
                                             m_randomTexture.GetAddressOf()));
    D3D11_SHADER_RESOURCE_VIEW_DESC randomTexSrvDesc{
        .Format = DXGI_FORMAT_R32G32B32A32_FLOAT,
        .ViewDimension = D3D11_SRV_DIMENSION_TEXTURE1D,
        .Texture1D = D3D11_TEX1D_SRV{0, 1}};
    DX::ThrowIfFailed(device.CreateShaderResourceView(
        m_randomTexture.Get(), &randomTexSrvDesc,
        m_randomTexSrv.GetAddressOf()));
}

void ParticleSystem::CreateInitialEmitter(ID3D11Device& device)
{
    const D3D11_BUFFER_DESC emitterBufferDesc{.ByteWidth = sizeof(Particle),
                                              .Usage = D3D11_USAGE_DEFAULT,
                                              .BindFlags =
                                                  D3D11_BIND_VERTEX_BUFFER,
                                              .CPUAccessFlags = 0};
    Particle emitter{.initialPos = {},
                     .intialVel = {},
                     .size = {},
                     .currentAge = 0.0f,
                     .type = static_cast<std::uint32_t>(ParticleType::kEmit)};

    D3D11_SUBRESOURCE_DATA subResData = {.pSysMem = &emitter};
    DX::ThrowIfFailed(device.CreateBuffer(&emitterBufferDesc, &subResData,
                                          m_initialParticle.GetAddressOf()));
}

void ParticleSystem::CreateParticleBuffers(ID3D11Device& device)
{
    const D3D11_BUFFER_DESC outputBufferDesc{
        .ByteWidth = sizeof(Particle) * m_maxParticles,
        .Usage = D3D11_USAGE_DEFAULT,
        .BindFlags = D3D11_BIND_VERTEX_BUFFER | D3D11_BIND_STREAM_OUTPUT,
        .CPUAccessFlags = 0};
    DX::ThrowIfFailed(device.CreateBuffer(&outputBufferDesc, nullptr,
                                          m_particleBuffer1.GetAddressOf()));

    DX::ThrowIfFailed(device.CreateBuffer(&outputBufferDesc, nullptr,
                                          m_particleBuffer2.GetAddressOf()));
}

void ParticleSystem::CreateShaders(ID3D11Device& device)
{
    const UINT bufferStrides[] = { sizeof(Particle) };
    D3D11_SO_DECLARATION_ENTRY entries[] = {
        {0, "POSITION", 0, 0, 4, 0}, 
        {0, "VELOCITY", 0, 0, 4, 0},
        {0, "SIZE", 0, 0, 2, 0},     {0, "LIFETIME", 0, 0, 1, 0},
        {0, "TYPE", 0, 0, 1, 0}};
    // geometry shaders
    ThrowIfFailed(device.CreateGeometryShaderWithStreamOutput(
        g_ParticleGenerationGSByteCode,
        std::size(g_ParticleGenerationGSByteCode), entries, std::size(entries),
        bufferStrides, std::size(bufferStrides), D3D11_SO_NO_RASTERIZED_STREAM,
        nullptr, m_particleGenerationShader.GetAddressOf()));
    ThrowIfFailed(device.CreateGeometryShader(
        g_ParticleExpansionGSByteCode, std::size(g_ParticleExpansionGSByteCode),
        nullptr, m_particleExpandingShader.GetAddressOf()));
    // vertex shaders
    ThrowIfFailed(device.CreateVertexShader(
        g_ForwardingVSByteCode, std::size(g_ForwardingVSByteCode), nullptr,
        m_forwardingVS.GetAddressOf()));
    ThrowIfFailed(device.CreateVertexShader(
        g_ParticleRenderingVSByteCode, std::size(g_ParticleRenderingVSByteCode),
        nullptr, m_expandingVS.GetAddressOf()));
    // pixel shaders
    ThrowIfFailed(device.CreatePixelShader(
        g_ParticleRenderingPSByteCode, std::size(g_ParticleRenderingPSByteCode),
        nullptr, m_samplingPS.GetAddressOf()));
}

void ParticleSystem::CreateInputLayouts(ID3D11Device& device)
{
    D3D11_INPUT_ELEMENT_DESC particleElementDesc[] = {
        {"POSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0,
         D3D11_APPEND_ALIGNED_ELEMENT},
        {"VELOCITY", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0,
         D3D11_APPEND_ALIGNED_ELEMENT},
        {"SIZE", 0, DXGI_FORMAT_R32G32_FLOAT, 0, D3D11_APPEND_ALIGNED_ELEMENT},
        {"LIFETIME", 0, DXGI_FORMAT_R32_FLOAT, 0, 40},
        {"TYPE", 0, DXGI_FORMAT_R32_UINT, 0, 44},
    };
    ThrowIfFailed(device.CreateInputLayout(
        particleElementDesc, std::size(particleElementDesc),
        g_ForwardingVSByteCode, std::size(g_ForwardingVSByteCode),
        m_particleInputLayout.GetAddressOf()));
}


template<typename T>
void MakeConstantBuffer(ID3D11Device& device3D, wrl::ComPtr<ID3D11Buffer>& outBuffer)
{
    D3D11_BUFFER_DESC cbDesc = {
        .ByteWidth = static_cast<UINT>(sizeof(T)),
        .Usage = D3D11_USAGE_DEFAULT,
        .BindFlags = D3D11_BIND_CONSTANT_BUFFER,
        //.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE,
    };
    ThrowIfFailed(device3D.CreateBuffer(&cbDesc, nullptr, outBuffer.GetAddressOf()));
}

template<typename T>
void UpdateConstantBuffer(ID3D11DeviceContext& context3D,
    ID3D11Buffer* buffer, const T* data)
{
    context3D.UpdateSubresource(buffer,
        0, nullptr, data, 0, 0);
}

void ParticleSystem::CreateConstantBuffers(ID3D11Device& device)
{
    MakeConstantBuffer<ParticleExpansionGSCbView>(device, m_particleExpansionGSViewCb);
    MakeConstantBuffer<ParticleGenGSCbPerFrame>(device, m_particleGenerationGSPerFrameCb);
    MakeConstantBuffer<ParticleGenGSCbFixed>(device, m_particleGenerationFixedCb);
    MakeConstantBuffer<ParticleRenderingVSCbFixed>(device, m_particleRenderingVSFixedCb);
}

void ParticleSystem::UpdateConstantBuffers(ID3D11DeviceContext& context3D, const Camera& camera)
{
    ParticleGenGSCbPerFrame particleGenGSCbPerFrame{
        .initialWorldPos = m_position,
        .timeElapsed = m_timeElapsed,
        .gameTime = m_now
    };

    UpdateConstantBuffer(context3D, m_particleGenerationGSPerFrameCb.Get(), &particleGenGSCbPerFrame);
    
    ParticleGenGSCbFixed particleGenGSCbFixed{
        .startLifetime = m_age
    };
    
    UpdateConstantBuffer(context3D, m_particleGenerationFixedCb.Get(), &particleGenGSCbFixed);

    ParticleExpansionGSCbView cbView{
        .viewProjMatrix = camera.viewMatrix * camera.projMatrix,
        .eyePos = camera.eyePos
    };

    UpdateConstantBuffer(context3D, m_particleExpansionGSViewCb.Get(), &cbView);

    ParticleRenderingVSCbFixed renderingVSCb{
        .accel = {0.0f, 7.8f, 0.0f},
        .startLifetime = m_age
    };

    UpdateConstantBuffer(context3D, m_particleRenderingVSFixedCb.Get(), &renderingVSCb);
}

void ParticleSystem::CreateSamplers(ID3D11Device& device)
{
    CD3D11_SAMPLER_DESC linearWrapDesc{ D3D11_DEFAULT };
    linearWrapDesc.AddressU = D3D11_TEXTURE_ADDRESS_WRAP;
    linearWrapDesc.AddressV = D3D11_TEXTURE_ADDRESS_WRAP;
    ThrowIfFailed(device.CreateSamplerState(&linearWrapDesc, m_linearWrapSampler.GetAddressOf()));

    linearWrapDesc.AddressU = D3D11_TEXTURE_ADDRESS_CLAMP;
    linearWrapDesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
    ThrowIfFailed(device.CreateSamplerState(&linearWrapDesc, m_linearClampSampler.GetAddressOf()));

    D3D11_DEPTH_STENCIL_DESC dsDesc{
        .DepthEnable = FALSE,
        .StencilEnable = FALSE
    };

    ThrowIfFailed(device.CreateDepthStencilState(&dsDesc, m_disableDepthStencilState.GetAddressOf()));

}

void ParticleSystem::SetupGenerationGS(ID3D11DeviceContext& context3D)
{
    context3D.GSSetShader(m_particleGenerationShader.Get(), nullptr, 0);
    ID3D11SamplerState* particleGenGSSamplers[] = { m_linearWrapSampler.Get() };
    context3D.GSSetSamplers(0, 1, particleGenGSSamplers);
        ID3D11Buffer* generationGSCbs[] = { m_particleGenerationGSPerFrameCb.Get(), m_particleGenerationFixedCb.Get() };
    context3D.GSSetConstantBuffers(0, std::size(generationGSCbs), generationGSCbs);
    ID3D11ShaderResourceView* srvs[] = { m_randomTexSrv.Get() };
    context3D.GSSetShaderResources(0, 1, srvs);
}

void ParticleSystem::SetupExpansionGS(ID3D11DeviceContext& context3D)
{
    context3D.GSSetShader(m_particleExpandingShader.Get(), nullptr, 0);
    ID3D11Buffer* cbs[] = { m_particleExpansionGSViewCb.Get() };
    D3D11_BUFFER_DESC desc;
    m_particleExpansionGSViewCb->GetDesc(&desc);
    context3D.GSSetConstantBuffers(0, 1, cbs);
}

void ParticleSystem::SetupRenderingPS(ID3D11DeviceContext& context3D)
{
    ID3D11ShaderResourceView* srvs[] = { m_billboard.Get() };
    context3D.PSSetShaderResources(0, 1, srvs);
    ID3D11SamplerState* samplers[] = { m_linearClampSampler.Get() };
    context3D.PSSetSamplers(0, 1, samplers);
    context3D.PSSetShader(m_samplingPS.Get(), nullptr, 0);
}

void ParticleSystem::SetupRenderingVS(ID3D11DeviceContext& context3D)
{
    ID3D11Buffer* cbs[] = { m_particleRenderingVSFixedCb.Get() };
    context3D.VSSetConstantBuffers(0, 1, cbs);
    context3D.VSSetShader(m_expandingVS.Get(), nullptr, 0);
}

void ParticleSystem::CreateAlphaBlendingState(ID3D11Device& device)
{
    D3D11_RENDER_TARGET_BLEND_DESC rtBlendingDesc{
        .BlendEnable = TRUE,
        .SrcBlend = D3D11_BLEND_SRC_ALPHA,
        .DestBlend = D3D11_BLEND_INV_SRC_ALPHA,
        .BlendOp = D3D11_BLEND_OP_ADD,
        .SrcBlendAlpha = D3D11_BLEND_ONE,
        .DestBlendAlpha = D3D11_BLEND_ZERO,
        .BlendOpAlpha = D3D11_BLEND_OP_ADD,
        .RenderTargetWriteMask = D3D11_COLOR_WRITE_ENABLE_ALL
    };

    D3D11_BLEND_DESC blendingDesc{
        .AlphaToCoverageEnable = FALSE,
        .IndependentBlendEnable = FALSE,
        .RenderTarget {rtBlendingDesc}
    };
    DX::ThrowIfFailed(device.CreateBlendState(&blendingDesc, m_alphaBlending.GetAddressOf()));

}

ParticleSystem::ParticleSystem(ID3D11Device& device3D, std::uint32_t maxParticles)
    : m_firstRun{ true }, m_maxParticles{1000}
{
    CreateRandomTexture(device3D);
    CreateInitialEmitter(device3D);
    CreateParticleBuffers(device3D);
    CreateShaders(device3D);
    CreateInputLayouts(device3D);
    CreateConstantBuffers(device3D);
    CreateSamplers(device3D);
    CreateAlphaBlendingState(device3D);
}

void ParticleSystem::SetMaxAge(float maxAge)
{
    m_age = maxAge;
}

void ParticleSystem::SetPosition(const DirectX::XMFLOAT3& position)
{
    m_position = position;
}

void ParticleSystem::SetBillboard(wrl::ComPtr<ID3D11ShaderResourceView> billboard)
{
    m_billboard = std::move(billboard);
}

void ParticleSystem::Update(float nowTime, float timeElapsed)
{
    m_timeElapsed = timeElapsed;
    m_now = nowTime;
}
